import os
import json
import logging
from datetime import datetime

class DataLogger:
    def __init__(self, output_dir='output'):
        self.output_dir = output_dir
        self.case_dir = None
        self.setup_logging()
        
    def setup_logging(self):
        """设置日志记录"""
        # 创建输出目录
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)
            
        # 创建本次计算的结果目录
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        self.case_dir = os.path.join(self.output_dir, f'case_{timestamp}')
        os.makedirs(self.case_dir)
        
        # 配置日志记录器
        log_file = os.path.join(self.case_dir, 'calculation.log')
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file),
                logging.StreamHandler()
            ]
        )
        
    def log_calculation_parameters(self, params):
        """记录计算参数"""
        params_file = os.path.join(self.case_dir, 'parameters.json')
        with open(params_file, 'w') as f:
            json.dump(params, f, indent=4)
        logging.info(f'Calculation parameters saved to {params_file}')
        
    def log_iteration_progress(self, iteration, mass_balance, energy_balance):
        """记录迭代进度"""
        logging.info(f'Iteration {iteration}:')
        logging.info(f'  Mass Balance Error: {mass_balance:.6e}')
        logging.info(f'  Energy Balance Error: {energy_balance:.6e}')
        
    def log_convergence_status(self, converged, total_iterations):
        """记录收敛状态"""
        if converged:
            logging.info(f'Calculation converged after {total_iterations} iterations')
        else:
            logging.warning(f'Calculation failed to converge after {total_iterations} iterations')
            
    def save_results(self, results_dict):
        """保存计算结果"""
        results_file = os.path.join(self.case_dir, 'results.json')
        with open(results_file, 'w') as f:
            json.dump(results_dict, f, indent=4)
        logging.info(f'Results saved to {results_file}')
        
    def generate_calculation_report(self, subchannel, convergence_checker):
        """生成计算报告"""
        report = {
            'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
            'convergence_status': {
                'mass_balance': convergence_checker.mass_balance_history[-1],
                'energy_balance': convergence_checker.energy_balance_history[-1]
            },
            'final_results': {
                'average_void_fraction': float(subchannel.void_fraction.mean()),
                'average_pressure': float(subchannel.pressure.mean()),
                'max_void_fraction': float(subchannel.void_fraction.max()),
                'min_pressure': float(subchannel.pressure.min())
            }
        }
        
        report_file = os.path.join(self.case_dir, 'calculation_report.json')
        with open(report_file, 'w') as f:
            json.dump(report, f, indent=4)
        logging.info(f'Calculation report generated: {report_file}') 